//////////////////////////////////////////////////////////////////////////
//////////////////           moonlander.cxx      /////////////////////////
//////////////////////////////////////////////////////////////////////////
////////////////           PSOPT  Example             ////////////////////
//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
//////// Title:         Moonlander problem                ////////////////
//////// Last modified: 05 January 2009                   ////////////////
//////// Reference:     PROPT Handbook                    ////////////////
//////// (See PSOPT handbook for full reference)          ////////////////
//////////////////////////////////////////////////////////////////////////
////////     Copyright (c) Victor M. Becerra, 2009        ////////////////
//////////////////////////////////////////////////////////////////////////
//////// This is part of the PSOPT software library, which ///////////////
//////// is distributed under the terms of the GNU Lesser ////////////////
//////// General Public License (LGPL)                    ////////////////
//////////////////////////////////////////////////////////////////////////

#include "psopt.h"


//////////////////////////////////////////////////////////////////////////
///////////////////  Define the end point (Mayer) cost function //////////
//////////////////////////////////////////////////////////////////////////

adouble endpoint_cost(adouble* initial_states, adouble* final_states,
                      adouble* parameters,adouble& t0, adouble& tf,
                      adouble* xad, int iphase, Workspace* workspace)
{
    return 0.0;
}

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the integrand (Lagrange) cost function  //////
//////////////////////////////////////////////////////////////////////////

adouble integrand_cost(adouble* states, adouble* controls, adouble* parameters,
                     adouble& time, adouble* xad, int iphase, Workspace* workspace)
{
    adouble thrust = controls[0];
    return  thrust;
}

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the DAE's ////////////////////////////////////
//////////////////////////////////////////////////////////////////////////

void dae(adouble* derivatives, adouble* path, adouble* states,
         adouble* controls, adouble* parameters, adouble& time,
         adouble* xad, int iphase, Workspace* workspace)
{
   adouble altitude_dot, speed_dot, mass_dot;

   adouble altitude  = states[0];
   adouble speed     = states[1];
   adouble mass      = states[2];

   adouble thrust = controls[0];

   double exhaust_velocity = 2.349;
   double gravity          = 1.0;

   altitude_dot =  speed;
   speed_dot    = -gravity + thrust/mass;
   mass_dot     = -thrust/exhaust_velocity;

   derivatives[0] = altitude_dot;
   derivatives[1] = speed_dot;
   derivatives[2] = mass_dot;
}

////////////////////////////////////////////////////////////////////////////
///////////////////  Define the events function ////////////////////////////
////////////////////////////////////////////////////////////////////////////

void events(adouble* e, adouble* initial_states, adouble* final_states,
            adouble* parameters,adouble& t0, adouble& tf, adouble* xad,
            int iphase, Workspace* workspace)
{
   adouble altitude_i = initial_states[0];
   adouble speed_i    = initial_states[1];
   adouble mass_i     = initial_states[2];
   adouble altitude_f = final_states[0];
   adouble speed_f    = final_states[1];

   e[0] = altitude_i;
   e[1] = speed_i;
   e[2] = mass_i;
   e[3] = altitude_f;
   e[4] = speed_f;

}

///////////////////////////////////////////////////////////////////////////
///////////////////  Define the phase linkages function ///////////////////
///////////////////////////////////////////////////////////////////////////

void linkages( adouble* linkages, adouble* xad, Workspace* workspace)
{
  // No linkages as this is a single phase problem
}


////////////////////////////////////////////////////////////////////////////
///////////////////  Define the main routine ///////////////////////////////
////////////////////////////////////////////////////////////////////////////

int main(void)
{
////////////////////////////////////////////////////////////////////////////
///////////////////  Declare key structures ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    Alg  algorithm;
    Sol  solution;
    Prob problem;

////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem name  ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    problem.name        		= "Moon Lander Problem";
    problem.outfilename                 = "moon.txt";

////////////////////////////////////////////////////////////////////////////
////////////  Define problem level constants & do level 1 setup ////////////
////////////////////////////////////////////////////////////////////////////

    problem.nphases   			= 1;
    problem.nlinkages                   = 0;

    psopt_level1_setup(problem);

/////////////////////////////////////////////////////////////////////////////
/////////   Define phase related information & do level 2 setup  ////////////
/////////////////////////////////////////////////////////////////////////////

    problem.phases(1).nstates   		= 3;
    problem.phases(1).ncontrols 		= 1;
    problem.phases(1).nevents   		= 5;
    problem.phases(1).npath     		= 0;
    problem.phases(1).nodes                     = 70;

    psopt_level2_setup(problem, algorithm);

    int nodes = 70;

////////////////////////////////////////////////////////////////////////////
///////////////////  Declare DMatrix objects to store results //////////////
////////////////////////////////////////////////////////////////////////////

    DMatrix x, u, t;
    DMatrix lambda, H;

////////////////////////////////////////////////////////////////////////////
///////////////////  Enter problem bounds information //////////////////////
////////////////////////////////////////////////////////////////////////////


    double altitudeL    = -20.0;
    double speedL 	= -20.0;
    double massL 	= 0.01;
    double altitudeU 	= 20.0;
    double speedU 	= 20.0;
    double massU 	= 1.0;

    double thrustL      = 0.0;
    double thrustU      = 1.227;

    double altitude_i  	= 1.0;
    double speed_i      = -0.783;
    double mass_i       = 1.0;

    double altitude_f   = 0.0;
    double speed_f      = 0.0;


    problem.phases(1).bounds.lower.states(1) = altitudeL;
    problem.phases(1).bounds.lower.states(2) = speedL;
    problem.phases(1).bounds.lower.states(3) = massL;

    problem.phases(1).bounds.upper.states(1) = altitudeU;
    problem.phases(1).bounds.upper.states(2) = speedU;
    problem.phases(1).bounds.upper.states(3) = massU;

    problem.phases(1).bounds.lower.controls(1) = thrustL;
    problem.phases(1).bounds.upper.controls(1) = thrustU;

    problem.phases(1).bounds.lower.events(1) = altitude_i;
    problem.phases(1).bounds.lower.events(2) = speed_i;
    problem.phases(1).bounds.lower.events(3) = mass_i;
    problem.phases(1).bounds.lower.events(4) = altitude_f;
    problem.phases(1).bounds.lower.events(5) = speed_f;

    problem.phases(1).bounds.upper.events =  problem.phases(1).bounds.lower.events;

    problem.phases(1).bounds.lower.StartTime    = 0.0;
    problem.phases(1).bounds.upper.StartTime    = 0.0;

    problem.phases(1).bounds.lower.EndTime      = 0.0;
    problem.phases(1).bounds.upper.EndTime      = 1000.0;


////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem functions  ///////////////////////////
////////////////////////////////////////////////////////////////////////////


    problem.integrand_cost 	= &integrand_cost;
    problem.endpoint_cost 	= &endpoint_cost;
    problem.dae 		= &dae;
    problem.events 		= &events;
    problem.linkages		= &linkages;

////////////////////////////////////////////////////////////////////////////
///////////////////  Define & register initial guess ///////////////////////
////////////////////////////////////////////////////////////////////////////

    DMatrix states_guess(3,nodes+1);

    states_guess(1,colon()) = linspace(altitude_i, altitude_f, nodes+1);
    states_guess(2,colon()) = linspace(speed_i, speed_f, nodes+1);
    states_guess(3,colon()) = linspace(mass_i, massL, nodes+1);


    problem.phases(1).guess.controls = 0.5*(thrustL+thrustU)*ones(1,nodes+1);
    problem.phases(1).guess.states   = states_guess;
    problem.phases(1).guess.time     = linspace(0.0, 1.5, nodes+1);

////////////////////////////////////////////////////////////////////////////
///////////////////  Enter algorithm options  //////////////////////////////
////////////////////////////////////////////////////////////////////////////

    algorithm.nlp_iter_max 		 = 1000;
    algorithm.nlp_tolerance   		 = 1.e-6;
    algorithm.nlp_method  		 = "IPOPT";
    algorithm.scaling     		 = "automatic";
    algorithm.derivatives 		 = "automatic";

////////////////////////////////////////////////////////////////////////////
///////////////////  Now call PSOPT to solve the problem   //////////////////
////////////////////////////////////////////////////////////////////////////

    psopt(solution, problem, algorithm);

////////////////////////////////////////////////////////////////////////////
///////////  Extract relevant variables from solution structure   //////////
////////////////////////////////////////////////////////////////////////////



    x 		= solution.get_states_in_phase(1);
    u 		= solution.get_controls_in_phase(1);
    t 		= solution.get_time_in_phase(1);
    lambda 	= solution.get_dual_costates_in_phase(1);
    H 		= solution.get_dual_hamiltonian_in_phase(1);

////////////////////////////////////////////////////////////////////////////
///////////  Save solution data to files if desired ////////////////////////
////////////////////////////////////////////////////////////////////////////

    x.Save("x.dat");
    u.Save("u.dat");
    t.Save("t.dat");
    lambda.Save("lambda.dat");
    H.Save("H.dat");

////////////////////////////////////////////////////////////////////////////
///////////  Plot some results if desired (requires gnuplot) ///////////////
////////////////////////////////////////////////////////////////////////////

    plot(t,x,problem.name, "time (s)", "states", "altitude speed mass");

    plot(t,u,problem.name, "time (s)", "control", "thrust");

    plot(t,x,problem.name, "time (s)", "states", "altitude speed mass",
                           "pdf", "moon_states.pdf");

    plot(t,u,problem.name, "time (s)", "control", "thrust",
                           "pdf", "moon_control.pdf");


}

////////////////////////////////////////////////////////////////////////////
///////////////////////      END OF FILE     ///////////////////////////////
////////////////////////////////////////////////////////////////////////////

